import os
import glob
import argparse
from pathlib import Path
import gmalg

def sm9_dec(mpk_e_file, msk_e_file, sk_e_file, uid_file, plain_file, cipher_file):
    hid_e = b"\x03"
    mpk_e_temp = read_binary_file_(mpk_e_file)
    mpk_e = bytes.fromhex("04") + mpk_e_temp
    msk_e = read_binary_file_(msk_e_file)
    sk_e_temp = read_binary_file_(sk_e_file)
    sk_e = bytes.fromhex("04") + sk_e_temp
    uid  = read_binary_file_(uid_file)
    plain  = read_binary_file_(plain_file)
    cipher_temp = read_binary_file_(cipher_file)
    cipher = bytes.fromhex("04") + cipher_temp
    #print("id=", uid.hex())
    #print("mpk_e_temp=", mpk_e_temp.hex())
    #print("msk_e=", msk_e.hex())
    #print("ske_e_temp=", sk_e_temp.hex())
    #print("plain=", plain.hex())
    #print("cipher_temp=", cipher_temp.hex())
    #print(cipher.hex())

    kgc = gmalg.SM9KGC(hid_e=hid_e, msk_e=msk_e, mpk_e=mpk_e)
    sk_e_creat  = kgc.generate_sk_encrypt(uid)
    assert sk_e_creat == sk_e, "验证失败"
    sm9 = gmalg.SM9(hid_e=hid_e, mpk_e=mpk_e, sk_e=sk_e, uid=uid)
    assert sm9.decrypt(cipher) == plain, "验证失败"

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['enc', 'dec'], default='enc', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode
    print (f"模式:{mode}\n")

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    if mode == 'enc' :
        for i in range(20):
            # 读取文件内容
            r_file = os.path.join(directory, f'r_{i}')
            id_file = os.path.join(directory, f'id')
            Ppub_e_file = os.path.join(directory, f'Ppub_e')
            ke_file = os.path.join(directory, f'ke')
            de_file = os.path.join(directory, f'de')
            plain_file = os.path.join(directory, f'msg')
            cipher_file = os.path.join(directory, f'cipher_{i}')

            sm9_dec(Ppub_e_file, ke_file, de_file, id_file, plain_file, cipher_file)
            r_content = read_binary_file(r_file)
            id_content = read_binary_file(id_file)
            Ppub_e_content = read_binary_file(Ppub_e_file)
            plain_content = read_binary_file(plain_file)
            cipher_content = read_binary_file(cipher_file)

            # 添加到合并内容
            merged_content += f"随机数=\n{r_content}\n"
            merged_content += f"用户标识=\n{id_content}\n"
            merged_content += f"公钥=\n{Ppub_e_content}\n"
            merged_content += f"加密方式= 0\n"
            merged_content += f"明文=\n{plain_content}\n"
            merged_content += f"密文=\n{cipher_content}\n"
            merged_content += "\n"
    else: 
        id_file = os.path.join(directory, f'id')
        Ppub_e_file = os.path.join(directory, f'Ppub_e')
        ke_file = os.path.join(directory, f'ke')
        de_file = os.path.join(directory, f'de')
        cipher_file = os.path.join(directory, f'cipher')
        replain_file = os.path.join(directory, f'replain')

        sm9_dec(Ppub_e_file, ke_file, de_file, id_file, replain_file, cipher_file)
        id_content = read_binary_file(id_file)
        de_content = read_binary_file(de_file)
        cipher_content = read_binary_file(cipher_file)
        replain_content = read_binary_file(replain_file)

        # 添加到合并内容
        merged_content += f"用户私钥=\n{de_content}\n"
        merged_content += f"用户标识=\n{id_content}\n"
        merged_content += f"加密方式= 0\n"
        merged_content += f"密文=\n{cipher_content}\n"
        merged_content += f"明文=\n{replain_content}\n"
        merged_content += "\n"

    # 写入到txt文件
    if mode == 'enc':
        output_file = 'sm9加密性能.txt'
    else:
        output_file = 'sm9解密性能.txt'

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
